import numpy as np
import matplotlib.pyplot as plt
import matplotlib.ticker as ticker

plt.rcParams.update({
    'font.size': 20,
    'font.family': 'serif',
    'font.serif': ['Times New Roman'],
    'svg.fonttype': 'none'
})


def plot_it(swap_vals, greedy_vals, optimal_vals, idx_ratio_1, optimal_ratio, optimal_ratio_1, x_labels, colors, border_colors, fig_name):
    colors = ['#4285f4', '#34a853', '#fbbc05']  # 分别对应 swap, greedy, optimal 的颜色
    border_colors = ["#002257", "#003a10", "#523d00"] # 分别对应 swap, greedy, optimal 的颜色
    
    swap_vals[:, 1] = swap_vals[:, 1] * np.array(idx_ratio_1)
    greedy_vals[:, 1] = greedy_vals[:, 1] * np.array(idx_ratio_1)
    optimal_vals[:, 0] = optimal_vals[:, 0] * np.array(optimal_ratio)
    optimal_vals[:, 1] = optimal_vals[:, 1] * \
        np.array(idx_ratio_1) * np.array(optimal_ratio_1)
    x = np.arange(len(x_labels))
    bar_width = 0.2
    offsets = np.array([-1.25, 0, 1.25]) * bar_width
    
    # 创建第一个图表（Execution Time）
    fig1, ax1 = plt.subplots(figsize=(8, 5))
    scheme_names = ['SO', 'MECH', 'JointST']
    for i, vals in enumerate([swap_vals[:, 0], greedy_vals[:, 0], optimal_vals[:, 0]]):
        ax1.bar(x + offsets[i], vals, width=bar_width,
                color=colors[i], alpha=1.0, label=scheme_names[i], edgecolor=border_colors[i], linewidth=1)
        ax1.plot(x + offsets[i], vals, color=colors[i],
                 marker='o', markersize=8, linestyle='-', linewidth=2)
    
    ax1.set_ylabel('Task completion time', fontsize=24)
    ax1.set_xlabel('Number of computation qubits', fontsize=24)
    ax1.set_xticks(x)
    ax1.set_xticklabels(x_labels)
    
    # 设置Y轴为科学计数法格式
    ax1.yaxis.set_major_formatter(ticker.ScalarFormatter(useMathText=True))
    ax1.yaxis.get_major_formatter().set_scientific(True)
    ax1.yaxis.get_major_formatter().set_powerlimits((0, 0))
    ax1.tick_params(axis='both', which='major', labelsize=24)
    ax1.legend(loc='upper left', fontsize=18)
    
    # 保存第一个图表
    plt.tight_layout()
    plt.savefig(fig_name.replace('.pdf', '_time.pdf'), format='pdf')
    plt.close(fig1)
    
    # 创建第二个图表（Avg. Ancilla Num）
    fig2, ax2 = plt.subplots(figsize=(8, 5))
    for i, vals in enumerate([greedy_vals[:, 1], optimal_vals[:, 1]]):
        # 注意：这里需要调整索引，因为没有swap方案的数据
        ax2.bar(x + offsets[i+1], vals, width=bar_width,
                color=colors[i+1], alpha=1.0, label=scheme_names[i+1], edgecolor=border_colors[i+1], linewidth=1)
        ax2.plot(x + offsets[i+1], vals, color=colors[i+1],
                 marker='o', markersize=8, linestyle='-', linewidth=2)
    
    ax2.set_ylabel('Avg. auxiliary qubit num', fontsize=24)
    ax2.set_xlabel('Number of computation qubits', fontsize=24)
    ax2.set_xticks(x)
    ax2.set_xticklabels(x_labels)
    
    # 设置Y轴为科学计数法格式
    ax2.yaxis.set_major_formatter(ticker.ScalarFormatter(useMathText=True))
    ax2.yaxis.get_major_formatter().set_scientific(True)
    ax2.yaxis.get_major_formatter().set_powerlimits((0, 0))
    ax2.tick_params(axis='both', which='major', labelsize=24)
    ax2.legend(loc='upper left', fontsize=18)
    
    # 保存第二个图表
    plt.tight_layout()
    plt.savefig(fig_name.replace('.pdf', '_ancilla.pdf'), format='pdf')
    plt.close(fig2)